{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BentoML Integration Guide"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[BentoML](https://github.com/bentoml/BentoML/) is an open-source framework designed for building,\n",
    "shipping, and scaling AI applications. It allows users to easily package and serve diffusion models\n",
    "for production, ensuring reliable and efficient deployments. It features out-of-the-box operational\n",
    "management tools like monitoring and tracing, and facilitates the deployment to various cloud platforms\n",
    "with ease. BentoML's distributed architecture and the separation of API server logic from\n",
    "model inference logic enable efficient scaling of deployments, even with budget constraints.\n",
    "As a result, integrating it with Diffusers provides a valuable tool for real-world deployments.\n",
    "\n",
    "This tutorial demonstrates how to integrate BentoML with Diffusers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Install [Diffusers](https://huggingface.co/docs/diffusers/installation).\n",
    "- Install BentoML by running `pip install bentoml`. For more information, see the [BentoML documentation](https://docs.bentoml.com)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import a diffusion model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, you need to prepare the model. BentoML has its own [Model Store](https://docs.bentoml.com/en/latest/concepts/model.html)\n",
    "for model management. Create a `download_model.py` file as below to import a diffusion model into BentoML's Model\n",
    "Store:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import bentoml\n",
    "\n",
    "bentoml.diffusers.import_model(\n",
    "    \"sd2.1\",  # Model tag in the BentoML Model Store\n",
    "    \"stabilityai/stable-diffusion-2-1\",  # Hugging Face model identifier\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code snippet downloads the Stable Diffusion 2.1 model (using it's repo id\n",
    "`stabilityai/stable-diffusion-2-1`) from the Hugging Face Hub (or use the cached download\n",
    "files if the model is already downloaded) and imports it into the BentoML Model\n",
    "Store with the name `sd2.1`.\n",
    "\n",
    "For models already fine-tuned and stored on disk, you can provide the path instead of\n",
    "the repo id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import bentoml\n",
    "\n",
    "bentoml.diffusers.import_model(\n",
    "    \"sd2.1-local\",\n",
    "    \"./local_stable_diffusion_2.1/\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can view the model in the Model Store:\n",
    "\n",
    "```\n",
    "bentoml models list\n",
    "\n",
    "Tag                                                                 Module                              Size       Creation Time       \n",
    "sd2.1:ysrlmubascajwnry                                              bentoml.diffusers                   33.85 GiB  2023-07-12 16:47:44 \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Turn a diffusion model into a RESTful service with BentoML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the diffusion model is in BentoML's Model Store, you can implement a text-to-image\n",
    "service with it. The Stable Diffusion model accepts various arguments\n",
    "in addition to the required prompt to guide the image generation process.\n",
    "To validate these input arguments, use BentoML's [pydantic](https://github.com/pydantic/pydantic) integration.\n",
    "Create a `sdargs.py` file with an example pydantic model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing as t\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class SDArgs(BaseModel):\n",
    "    prompt: str\n",
    "    negative_prompt: t.Optional[str] = None\n",
    "    height: t.Optional[int] = 512\n",
    "    width: t.Optional[int] = 512\n",
    "\n",
    "    class Config:\n",
    "        extra = \"allow\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This pydantic model requires a string field `prompt` and three optional fields: `height`, `width`, and `negative_prompt`,\n",
    "each with corresponding types. The `extra = \"allow\"` line supports adding additional fields not defined in the `SDArgs` class.\n",
    "In a real-world scenario, you may define all the desired fields and not allow extra ones.\n",
    "\n",
    "Next, create a BentoML Service file that defines a Stable Diffusion service:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import bentoml\n",
    "from bentoml.io import Image, JSON\n",
    "\n",
    "from sdargs import SDArgs\n",
    "\n",
    "bento_model = bentoml.diffusers.get(\"sd2.1:latest\")\n",
    "sd21_runner = bento_model.to_runner(name=\"sd21-runner\")\n",
    "\n",
    "svc = bentoml.Service(\"stable-diffusion-21\", runners=[sd21_runner])\n",
    "\n",
    "\n",
    "@svc.api(input=JSON(pydantic_model=SDArgs), output=Image())\n",
    "async def txt2img(input_data):\n",
    "    kwargs = input_data.dict()\n",
    "    res = await sd21_runner.async_run(**kwargs)\n",
    "    images = res[0]\n",
    "    return images[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the file as `service.py`, and spin up a BentoML Service endpoint using:\n",
    "\n",
    "```\n",
    "bentoml serve service:svc\n",
    "```\n",
    "\n",
    "An HTTP server with `/txt2img` endpoint that accepts a JSON dictionary should be up at\n",
    "port 3000. Go to <http://127.0.0.1:3000> in your web browser to access the Swagger UI.\n",
    "\n",
    "You can also test the text-to-image generation using `curl` and write the returned image to\n",
    "`output.jpg`.\n",
    "\n",
    "```\n",
    "curl -X POST http://127.0.0.1:3000/txt2img \\\n",
    "     -H 'Content-Type: application/json' \\\n",
    "     -d \"{\\\"prompt\\\":\\\"a black cat\\\", \\\"height\\\":768, \\\"width\\\":768}\" \\\n",
    "     --output output.jpg\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Package a BentoML Service for cloud deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To deploy a BentoML Service, you need to pack it into a BentoML\n",
    "[Bento](https://docs.bentoml.com/en/latest/concepts/bento.html), a file archive with all the source code,\n",
    "models, data files, and dependencies. This can be done by providing a `bentofile.yaml` file as follows:\n",
    "\n",
    "```yaml\n",
    "service: \"service.py:svc\"\n",
    "include:\n",
    "  - \"service.py\"\n",
    "python:\n",
    "  packages:\n",
    "    - torch\n",
    "    - transformers\n",
    "    - accelerate\n",
    "    - diffusers\n",
    "    - triton\n",
    "    - xformers\n",
    "    - pydantic\n",
    "docker:\n",
    "    distro: debian\n",
    "    cuda_version: \"11.6\"\n",
    "```\n",
    "\n",
    "The `bentofile.yaml` file contains [Bento build\n",
    "options](https://docs.bentoml.com/en/latest/concepts/bento.html#bento-build-options),\n",
    "such as package dependencies and Docker options.\n",
    "\n",
    "Then you build a Bento using:\n",
    "\n",
    "```\n",
    "bentoml build\n",
    "```\n",
    "\n",
    "The output looks like:\n",
    "\n",
    "```\n",
    "Successfully built Bento(tag=\"stable-diffusion-21:crkuh7a7rw5bcasc\").\n",
    "\n",
    "Possible next steps:\n",
    "\n",
    " * Containerize your Bento with `bentoml containerize`:\n",
    "    $ bentoml containerize stable-diffusion-21:crkuh7a7rw5bcasc\n",
    "\n",
    " * Push to BentoCloud with `bentoml push`:\n",
    "    $ bentoml push stable-diffusion-21:crkuh7a7rw5bcasc\n",
    "```\n",
    "\n",
    "You can create a Docker image based on the Bento by running the following command and deploy it to a cloud provider.\n",
    "\n",
    "```\n",
    "bentoml containerize stable-diffusion-21:crkuh7a7rw5bcasc\n",
    "```\n",
    "\n",
    "If you want an end-to-end solution for deploying and managing models, you can push the Bento to [Yatai](https://github.com/bentoml/Yatai) or\n",
    "[BentoCloud](https://bentoml.com/cloud) for a distributed deployment.\n",
    "\n",
    "For more information about BentoML's integration with Diffusers, see the [BentoML Diffusers\n",
    "Guide](https://docs.bentoml.com/en/latest/frameworks/diffusers.html)."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
